#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
游戏截图信息提取工具
使用硅基流动的Qwen3-VL模型提取游戏截图中的区服、角色名和UID信息
"""

import sys
import os
import json
import base64
import requests
import subprocess
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, 
                             QWidget, QPushButton, QLabel, QLineEdit, QTextEdit, 
                             QTableWidget, QTableWidgetItem, QFileDialog, QMessageBox,
                             QProgressBar, QGroupBox, QGridLayout, QHeaderView,
                             QScrollArea, QFrame, QSpinBox, QComboBox, QMenu, QAction,
                             QDialog, QListWidget, QListWidgetItem)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QRect, QPoint
from PyQt5.QtGui import QRegion, QColor
from PyQt5.QtGui import QPixmap, QFont, QPainter, QPen, QBrush
from PIL import Image
import pandas as pd
from io import BytesIO
import re
import mss
import time
import cv2
import numpy as np
import webbrowser

class ConfigManager:
    """配置管理器"""
    
    def __init__(self):
        self.config_file = "config.json"
        self.config = self.load_config()
    
    def load_config(self):
        """加载配置文件"""
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"加载配置文件失败: {e}")
                return self.get_default_config()
        else:
            return self.get_default_config()
    
    def get_default_config(self):
        """获取默认配置"""
        return {
            "api_key": "",
            "save_path": os.getcwd(),
            "window_geometry": {
                "x": 100,
                "y": 100,
                "width": 750,
                "height": 600
            },
            "extraction_method": "未选择",
            "server_names": [
                "铭心之界", "拉普达", "轻云之月", "应许之地", "旷野之原", 
                "遗忘之境", "参天之林", "生命之树", "魔女小樱", "绽放之花", 
                "暗涌之泉", "缥缈之峰", "无尽之海"
            ]
        }
    
    def save_config(self):
        """保存配置文件"""
        try:
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存配置文件失败: {e}")
            return False
    
    def get(self, key, default=None):
        """获取配置值"""
        return self.config.get(key, default)
    
    def set(self, key, value):
        """设置配置值"""
        self.config[key] = value
        self.save_config()
    
    def get_api_key(self):
        """获取API密钥"""
        return self.get("api_key", "")
    
    def set_api_key(self, api_key):
        """设置API密钥"""
        self.set("api_key", api_key)
    
    def get_save_path(self):
        """获取保存路径"""
        return self.get("save_path", os.getcwd())
    
    def set_save_path(self, save_path):
        """设置保存路径"""
        self.set("save_path", save_path)
    
    def get_window_geometry(self):
        """获取窗口几何信息"""
        return self.get("window_geometry", {
            "x": 100, "y": 100, "width": 750, "height": 600
        })
    
    def set_window_geometry(self, geometry):
        """设置窗口几何信息"""
        self.set("window_geometry", geometry)
    
    def get_extraction_method(self):
        """获取提取方法"""
        return self.get("extraction_method", "未选择")
    
    def set_extraction_method(self, method):
        """设置提取方法"""
        self.set("extraction_method", method)
    
    def get_server_names(self):
        """获取区服名称列表"""
        return self.get("server_names", [
            "铭心之界", "拉普达", "轻云之月", "应许之地", "旷野之原", 
            "遗忘之境", "参天之林", "生命之树", "魔女小樱", "绽放之花", 
            "暗涌之泉", "缥缈之峰", "无尽之海"
        ])

class LoginDialog(QDialog):
    """登录对话框"""
    
    def __init__(self):
        super().__init__()
        self.password = "naikuai666"
        self.init_ui()
    
    def init_ui(self):
        """初始化登录界面"""
        self.setWindowTitle("角色信息提取工具V1.0.0 - 登录")
        self.setFixedSize(480, 380)
        self.setWindowFlags(Qt.WindowCloseButtonHint | Qt.WindowMinimizeButtonHint)
        
        # 设置窗口居中
        self.center_window()
        
        # 主布局
        layout = QVBoxLayout()
        layout.setSpacing(12)
        layout.setContentsMargins(25, 25, 25, 25)
        
        # 标题
        title_label = QLabel("角色信息提取工具V1.0.0")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setFont(QFont("Microsoft YaHei", 15, QFont.Bold))
        title_label.setStyleSheet("color: #2c3e50; margin-bottom: 2px; font-weight: bold;")
        layout.addWidget(title_label)
        
        # 版本信息
        version_label = QLabel("Version 1.0.0")
        version_label.setAlignment(Qt.AlignCenter)
        version_label.setFont(QFont("Microsoft YaHei", 8))
        version_label.setStyleSheet("color: #7f8c8d; margin-bottom: 8px;")
        layout.addWidget(version_label)
        
        # 副标题
        subtitle_label = QLabel("请输入密码以继续使用")
        subtitle_label.setAlignment(Qt.AlignCenter)
        subtitle_label.setFont(QFont("Microsoft YaHei", 10))
        subtitle_label.setStyleSheet("color: #34495e; margin-bottom: 8px; font-weight: normal;")
        layout.addWidget(subtitle_label)
        
        # 密码输入区域
        password_group = QGroupBox("登录验证")
        password_group.setFont(QFont("Microsoft YaHei", 9, QFont.Bold))
        password_group.setStyleSheet("""
            QGroupBox {
                font-weight: bold;
                border: 2px solid #bdc3c7;
                border-radius: 6px;
                margin-top: 5px;
                padding-top: 10px;
                background-color: #ffffff;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px 0 5px;
                color: #2c3e50;
            }
        """)
        password_layout = QVBoxLayout(password_group)
        password_layout.setSpacing(5)
        password_layout.setContentsMargins(10, 10, 10, 10)
        
        password_label = QLabel("密码:")
        password_label.setFont(QFont("Microsoft YaHei", 9, QFont.Bold))
        password_label.setStyleSheet("color: #2c3e50;")
        password_layout.addWidget(password_label)
        
        self.password_input = QLineEdit()
        self.password_input.setEchoMode(QLineEdit.Password)
        self.password_input.setPlaceholderText("请输入访问密码")
        self.password_input.setFont(QFont("Microsoft YaHei", 9))
        self.password_input.setStyleSheet("""
            QLineEdit {
                padding: 8px;
                border: 2px solid #bdc3c7;
                border-radius: 4px;
                font-size: 9px;
                background-color: #ffffff;
                color: #2c3e50;
                min-height: 20px;
            }
            QLineEdit:focus {
                border-color: #3498db;
                background-color: #f8f9fa;
            }
        """)
        self.password_input.returnPressed.connect(self.login)
        password_layout.addWidget(self.password_input)
        
        layout.addWidget(password_group)
        
        # 按钮布局
        button_layout = QHBoxLayout()
        button_layout.setSpacing(10)
        
        # 登录按钮
        self.login_btn = QPushButton("登录")
        self.login_btn.setFont(QFont("Microsoft YaHei", 10, QFont.Bold))
        self.login_btn.setStyleSheet("""
            QPushButton {
                background-color: #3498db;
                color: white;
                border: none;
                padding: 10px 20px;
                border-radius: 5px;
                font-weight: bold;
                min-width: 90px;
                font-size: 10px;
            }
            QPushButton:hover {
                background-color: #2980b9;
            }
            QPushButton:pressed {
                background-color: #21618c;
            }
        """)
        self.login_btn.clicked.connect(self.login)
        button_layout.addWidget(self.login_btn)
        
        # 退出按钮
        self.exit_btn = QPushButton("退出")
        self.exit_btn.setFont(QFont("Microsoft YaHei", 10, QFont.Bold))
        self.exit_btn.setStyleSheet("""
            QPushButton {
                background-color: #e74c3c;
                color: white;
                border: none;
                padding: 10px 20px;
                border-radius: 5px;
                font-weight: bold;
                min-width: 90px;
                font-size: 10px;
            }
            QPushButton:hover {
                background-color: #c0392b;
            }
            QPushButton:pressed {
                background-color: #a93226;
            }
        """)
        self.exit_btn.clicked.connect(self.exit_app)
        button_layout.addWidget(self.exit_btn)
        
        layout.addLayout(button_layout)
        
        # 联系信息
        contact_label = QLabel("联系作者: AJower@163.com  |  QQ: 3175268324")
        contact_label.setAlignment(Qt.AlignCenter)
        contact_label.setFont(QFont("Microsoft YaHei", 8))
        contact_label.setStyleSheet("color: #95a5a6; margin-top: 10px; padding: 6px; background-color: #f8f9fa; border-radius: 3px; font-weight: normal;")
        layout.addWidget(contact_label)
        
        self.setLayout(layout)
        
        # 设置样式
        self.setStyleSheet("""
            QDialog {
                background-color: #ecf0f1;
            }
        """)
    
    def center_window(self):
        """窗口居中"""
        screen = QApplication.desktop().screenGeometry()
        size = self.geometry()
        self.move(
            (screen.width() - size.width()) // 2,
            (screen.height() - size.height()) // 2
        )
    
    def login(self):
        """登录验证"""
        input_password = self.password_input.text()
        
        if input_password == self.password:
            self.accept()  # 登录成功
        else:
            QMessageBox.warning(self, "登录失败", "密码错误，请重新输入！")
            self.password_input.clear()
            self.password_input.setFocus()
    
    def exit_app(self):
        """退出应用"""
        sys.exit(0)

def check_dependencies():
    """检查依赖包是否已安装"""
    required_packages = [
        "PyQt5", "requests", "PIL", "pandas", 
        "openpyxl", "numpy", "mss", "cv2"
    ]
    
    missing_packages = []
    
    for package in required_packages:
        try:
            if package == "PIL":
                import PIL
            elif package == "cv2":
                import cv2
            else:
                __import__(package)
        except ImportError:
            missing_packages.append(package)
    
    return missing_packages

def install_dependencies():
    """安装依赖包"""
    print("正在检查并安装依赖包...")
    
    # 使用清华镜像源加速下载
    pip_install_cmd = [
        sys.executable, "-m", "pip", "install", 
        "-i", "https://pypi.tuna.tsinghua.edu.cn/simple",
        "--upgrade", "pip"
    ]
    
    # 依赖包列表
    packages = [
        "PyQt5==5.15.9",
        "requests==2.31.0", 
        "Pillow==10.0.0",
        "pandas==2.0.3",
        "openpyxl==3.1.2",
        "numpy==1.24.3",
        "mss==9.0.1",
        "opencv-python==4.8.1.78"
    ]
    
    try:
        # 升级pip
        print("正在升级pip...")
        subprocess.run(pip_install_cmd, check=True, capture_output=True)
        
        # 安装依赖包
        for package in packages:
            print(f"正在安装 {package}...")
            install_cmd = [
                sys.executable, "-m", "pip", "install",
                "-i", "https://pypi.tuna.tsinghua.edu.cn/simple",
                package
            ]
            subprocess.run(install_cmd, check=True, capture_output=True)
            print(f"✓ {package} 安装成功")
        
        print("所有依赖包安装完成！")
        return True
        
    except subprocess.CalledProcessError as e:
        print(f"安装依赖包失败: {e}")
        return False
    except Exception as e:
        print(f"安装过程中出现错误: {e}")
        return False

class ScreenshotWidget(QWidget):
    """OpenCV截图选择区域的小部件"""
    screenshot_taken = pyqtSignal(str)  # 截图完成信号，传递图片路径
    
    def __init__(self, save_path=""):
        super().__init__()
        self.start_pos = None
        self.end_pos = None
        self.drawing = False
        self.screenshot = None
        self.screenshot_pixmap = None
        self.save_path = save_path
        self.setWindowFlags(Qt.WindowStaysOnTopHint | Qt.FramelessWindowHint)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setStyleSheet("background-color: rgba(0, 0, 0, 150);")
        
        # 启用鼠标跟踪
        self.setMouseTracking(True)
        
    def start_screenshot(self):
        """开始截图"""
        try:
            print("开始获取全屏截图...")
            
            # 使用OpenCV获取全屏截图
            with mss.mss() as sct:
                monitors = sct.monitors
                print(f"检测到 {len(monitors)} 个显示器")
                monitor = monitors[1]  # 索引1是主显示器
                print(f"使用显示器: {monitor}")
                self.screenshot = sct.grab(monitor)
                print(f"截图尺寸: {self.screenshot.size}")
            
            # 转换为OpenCV格式
            img = Image.frombytes("RGB", self.screenshot.size, self.screenshot.bgra, "raw", "BGRX")
            self.cv_image = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            print(f"OpenCV图像尺寸: {self.cv_image.shape}")
            
            # 转换为QPixmap用于显示
            img_bytes = BytesIO()
            img.save(img_bytes, format='PNG')
            img_bytes.seek(0)
            
            self.screenshot_pixmap = QPixmap()
            success = self.screenshot_pixmap.loadFromData(img_bytes.getvalue())
            print(f"QPixmap加载: {'成功' if success else '失败'}")
            
            if not success:
                raise Exception("QPixmap加载失败")
            
            # 创建全屏窗口
            self.setGeometry(0, 0, img.width, img.height)
            print(f"设置窗口尺寸: {img.width} x {img.height}")
            
            # 显示窗口
            self.show()
            self.raise_()
            self.activateWindow()
            print("截图窗口已显示")
            
        except Exception as e:
            print(f"截图初始化失败: {e}")
            import traceback
            traceback.print_exc()
            self.close()
        
    def paintEvent(self, event):
        """绘制截图和选择区域"""
        painter = QPainter(self)
        
        # 绘制原始截图作为背景
        if self.screenshot_pixmap:
            painter.drawPixmap(self.rect(), self.screenshot_pixmap)
        
        # 只在有选择区域时绘制遮罩
        if self.start_pos and self.end_pos:
            rect = QRect(self.start_pos, self.end_pos).normalized()
            
            # 创建遮罩区域（整个屏幕减去选择区域）
            mask_region = QRegion(self.rect()).subtracted(QRegion(rect))
            painter.setClipRegion(mask_region)
            painter.fillRect(self.rect(), QBrush(Qt.black, Qt.SolidPattern))
            painter.setClipping(False)
            
            # 绘制选择区域（显示原始截图）
            if self.screenshot_pixmap:
                painter.drawPixmap(rect, self.screenshot_pixmap, rect)
            
            # 绘制选择框边框
            painter.setPen(QPen(Qt.red, 3))
            painter.setBrush(QBrush(Qt.transparent))
            painter.drawRect(rect)
            
            # 显示选择区域大小
            painter.setPen(QPen(Qt.white, 1))
            painter.setFont(QFont("Arial", 12, QFont.Bold))
            size_text = f"{rect.width()} x {rect.height()}"
            painter.drawText(rect.bottomRight() + QPoint(10, 10), size_text)
        else:
            # 初始状态，不绘制遮罩，只显示原始截图
            # 绘制十字准线
            if hasattr(self, 'current_mouse_pos'):
                pos = self.current_mouse_pos
                painter.setPen(QPen(Qt.white, 1, Qt.DashLine))
                
                # 绘制垂直线
                painter.drawLine(pos.x(), 0, pos.x(), self.height())
                # 绘制水平线
                painter.drawLine(0, pos.y(), self.width(), pos.y())
    
    def mousePressEvent(self, event):
        """鼠标按下事件"""
        if event.button() == Qt.LeftButton:
            self.start_pos = event.pos()
            self.end_pos = event.pos()
            self.drawing = True
            self.update()
    
    def mouseMoveEvent(self, event):
        """鼠标移动事件"""
        # 记录鼠标位置用于绘制十字准线
        self.current_mouse_pos = event.pos()
        
        if self.drawing:
            self.end_pos = event.pos()
            self.update()
        else:
            # 即使没有拖拽，也更新显示以绘制十字准线
            self.update()
    
    def mouseReleaseEvent(self, event):
        """鼠标释放事件"""
        if event.button() == Qt.LeftButton and self.drawing:
            self.end_pos = event.pos()
            self.drawing = False
            
            # 计算选择区域
            rect = QRect(self.start_pos, self.end_pos).normalized()
            
            # 如果选择区域太小，取消截图
            if rect.width() < 5 or rect.height() < 5:
                self.close()
                return
            
            # 截取选择区域
            self.capture_selected_area(rect)
    
    def keyPressEvent(self, event):
        """键盘事件"""
        if event.key() == Qt.Key_Escape:
            self.close()
    
    def capture_selected_area(self, rect):
        """截取选择的区域"""
        try:
            print(f"开始截取区域: {rect}")
            
            # 从OpenCV图像中截取选择区域
            x, y, w, h = rect.x(), rect.y(), rect.width(), rect.height()
            
            # 确保坐标在图像范围内
            x = max(0, min(x, self.cv_image.shape[1] - 1))
            y = max(0, min(y, self.cv_image.shape[0] - 1))
            w = min(w, self.cv_image.shape[1] - x)
            h = min(h, self.cv_image.shape[0] - y)
            
            print(f"调整后坐标: x={x}, y={y}, w={w}, h={h}")
            
            # 使用OpenCV截取区域
            cropped = self.cv_image[y:y+h, x:x+w]
            print(f"截取区域尺寸: {cropped.shape}")
            
            # 确保保存目录存在
            if self.save_path:
                os.makedirs(self.save_path, exist_ok=True)
                save_dir = self.save_path
            else:
                save_dir = os.getcwd()
            
            # 保存截图
            timestamp = int(time.time())
            filename = f"screenshot_{timestamp}.png"
            filepath = os.path.join(save_dir, filename)
            
            print(f"保存路径: {filepath}")
            
            # 使用OpenCV保存图像
            success = cv2.imwrite(filepath, cropped)
            
            if success and os.path.exists(filepath):
                file_size = os.path.getsize(filepath)
                print(f"截图已保存: {filepath} (大小: {file_size} 字节)")
                self.screenshot_taken.emit(filepath)
            else:
                raise Exception("OpenCV保存失败")
            
        except Exception as e:
            print(f"截图保存失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.close()

class AIExtractor(QThread):
    """AI信息提取线程"""
    progress_updated = pyqtSignal(int)
    result_ready = pyqtSignal(dict)
    error_occurred = pyqtSignal(str)
    
    def __init__(self, api_key, image_paths, server_names, extraction_method="unknown"):
        super().__init__()
        self.api_key = api_key
        self.image_paths = image_paths
        self.server_names = server_names
        self.extraction_method = extraction_method
        self.results = []
        
    def run(self):
        """运行AI提取任务"""
        total_images = len(self.image_paths)
        
        for i, image_path in enumerate(self.image_paths):
            try:
                # 更新进度
                self.progress_updated.emit(int((i / total_images) * 100))
                
                # 提取单张图片信息
                result = self.extract_image_info(image_path)
                self.results.append(result)
                
            except Exception as e:
                error_result = {
                    'image_path': image_path,
                    'server': '提取失败',
                    'character_name': '提取失败',
                    'uid': '提取失败',
                    'error': str(e)
                }
                self.results.append(error_result)
                
        # 完成所有提取
        self.progress_updated.emit(100)
        self.result_ready.emit({'results': self.results})
    
    def extract_image_info(self, image_path):
        """提取单张图片的信息"""
        try:
            # 读取并编码图片
            with open(image_path, 'rb') as f:
                image_data = f.read()
            
            # 转换为base64
            image_base64 = base64.b64encode(image_data).decode('utf-8')
            
            # 构建API请求
            url = "https://api.siliconflow.cn/v1/chat/completions"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            # 构建提示词
            server_list = "、".join(self.server_names)
            prompt = f"""请分析这张游戏截图，提取以下信息：
1. 区服名称（必须是以下之一：{server_list}）
   - 区服名称可能出现在用户名后面，如"@用户名 应许角色名:xxx"
   - 区服名称可能与"角色名"连在一起，如"应许角色名:xxx"
   - 如果图片中显示的是简称，请转换为完整名称
   - 例如：看到"应许"应识别为"应许之地"，看到"铭心"应识别为"铭心之界"
   - 支持的区服：{server_list}
2. 角色名称（通常在"角色名:"后面）
3. UID（数字，通常在"UID:"后面）

请以JSON格式返回结果，格式如下：
{{
    "server": "区服名称",
    "character_name": "角色名称", 
    "uid": "UID数字"
}}

如果无法识别某项信息，请返回"未识别"。"""
            
            data = {
                "model": "Qwen/Qwen3-VL-30B-A3B-Instruct",
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_base64}"
                                }
                            }
                        ]
                    }
                ],
                "max_tokens": 500,
                "temperature": 0.1
            }
            
            # 发送请求
            response = requests.post(url, headers=headers, json=data, timeout=30)
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            content = result['choices'][0]['message']['content']
            
            # 尝试解析JSON
            try:
                # 提取JSON部分
                json_match = re.search(r'\{.*\}', content, re.DOTALL)
                if json_match:
                    json_str = json_match.group()
                    extracted_data = json.loads(json_str)
                else:
                    # 如果无法找到JSON，尝试手动解析
                    extracted_data = self.parse_text_response(content)
                
                # 验证区服名称
                if extracted_data.get('server') not in self.server_names:
                    extracted_data['server'] = '未识别'
                
                # 验证UID是否为数字
                uid = extracted_data.get('uid', '')
                if not uid.isdigit():
                    extracted_data['uid'] = '未识别'
                
                return {
                    'image_path': os.path.basename(image_path),
                    'method': self.extraction_method,
                    'parse_type': 'AI解析',
                    'server': extracted_data.get('server', '未识别'),
                    'character_name': extracted_data.get('character_name', '未识别'),
                    'uid': extracted_data.get('uid', '未识别'),
                    'error': None
                }
                
            except (json.JSONDecodeError, KeyError) as e:
                # 如果JSON解析失败，返回原始内容
                return {
                    'image_path': os.path.basename(image_path),
                    'method': self.extraction_method,
                    'parse_type': '手动解析',
                    'server': '解析失败',
                    'character_name': '解析失败',
                    'uid': '解析失败',
                    'error': f"JSON解析失败: {str(e)}"
                }
                
        except Exception as e:
            return {
                'image_path': os.path.basename(image_path),
                'method': self.extraction_method,
                'parse_type': '提取失败',
                'server': '提取失败',
                'character_name': '提取失败',
                'uid': '提取失败',
                'error': str(e)
            }
    
    def parse_text_response(self, content):
        """手动解析文本响应"""
        result = {'server': '未识别', 'character_name': '未识别', 'uid': '未识别'}
        
        # 查找区服（支持简称匹配）
        # 首先尝试完整名称匹配
        for server in self.server_names:
            if server in content:
                result['server'] = server
                break
        
        # 如果没找到完整名称，尝试简称匹配
        if result['server'] == '未识别':
            server_shortcuts = {
                "铭心": "铭心之界", "拉普": "拉普达", "轻云": "轻云之月",
                "应许": "应许之地", "旷野": "旷野之原", "遗忘": "遗忘之境",
                "参天": "参天之林", "生命": "生命之树", "魔女": "魔女小樱",
                "绽放": "绽放之花", "暗涌": "暗涌之泉", "缥缈": "缥缈之峰",
                "无尽": "无尽之海", "界": "铭心之界", "达": "拉普达",
                "月": "轻云之月", "地": "应许之地", "原": "旷野之原",
                "境": "遗忘之境", "林": "参天之林", "树": "生命之树",
                "樱": "魔女小樱", "花": "绽放之花", "泉": "暗涌之泉",
                "峰": "缥缈之峰", "海": "无尽之海"
            }
            
            # 尝试匹配与"角色名"连在一起的格式，如"应许角色名:xxx"
            for shortcut, full_name in server_shortcuts.items():
                if f"{shortcut}角色名" in content or f"{shortcut} 角色名" in content:
                    result['server'] = full_name
                    break
            
            # 如果还没找到，尝试一般匹配
            if result['server'] == '未识别':
                for shortcut, full_name in server_shortcuts.items():
                    if shortcut in content:
                        result['server'] = full_name
                        break
        
        # 查找UID（数字）
        uid_match = re.search(r'uid["\']?\s*[:：]\s*["\']?(\d+)', content, re.IGNORECASE)
        if uid_match:
            result['uid'] = uid_match.group(1)
        
        # 查找角色名（在角色名关键词附近）
        char_patterns = [
            r'角色名["\']?\s*[:：]\s*["\']?([^"\',，\n]+)',
            r'character["\']?\s*[:：]\s*["\']?([^"\',，\n]+)',
            r'角色["\']?\s*[:：]\s*["\']?([^"\',，\n]+)'
        ]
        
        for pattern in char_patterns:
            char_match = re.search(pattern, content, re.IGNORECASE)
            if char_match:
                result['character_name'] = char_match.group(1).strip()
                break
        
        return result


class ScreenshotExtractor(QMainWindow):
    """主窗口类"""
    
    def __init__(self):
        super().__init__()
        # 初始化配置管理器
        self.config_manager = ConfigManager()
        
        # 从配置文件加载设置
        self.api_key = self.config_manager.get_api_key()
        self.save_path = self.config_manager.get_save_path()
        self.extraction_method = self.config_manager.get_extraction_method()
        self.server_names = self.config_manager.get_server_names()
        
        self.image_paths = []
        self.results = []
        self.screenshot_widget = None
        self.lottery_selected_results = []  # 存储抽取结果
        
        self.init_ui()
        
        # 加载窗口几何信息
        geometry = self.config_manager.get_window_geometry()
        self.setGeometry(geometry["x"], geometry["y"], geometry["width"], geometry["height"])
    
    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("角色信息提取工具V1.0.0")
        self.setGeometry(100, 100, 750, 600)  # 调整窗口尺寸以适应新功能
        
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 主布局 - 减小间距
        main_layout = QVBoxLayout(central_widget)
        main_layout.setSpacing(5)
        main_layout.setContentsMargins(8, 8, 8, 8)
        
        # 配置区域 - 合并API和路径设置
        config_group = QGroupBox("配置")
        config_layout = QGridLayout(config_group)
        config_layout.setSpacing(5)
        
        # API密钥
        config_layout.addWidget(QLabel("API密钥:"), 0, 0)
        self.api_key_input = QLineEdit()
        self.api_key_input.setPlaceholderText("请输入硅基流动API密钥")
        self.api_key_input.setEchoMode(QLineEdit.Password)
        config_layout.addWidget(self.api_key_input, 0, 1, 1, 2)
        
        # 保存路径
        config_layout.addWidget(QLabel("保存路径:"), 1, 0)
        self.save_path_input = QLineEdit()
        self.save_path_input.setText(self.save_path)
        self.save_path_input.setPlaceholderText("截图保存路径")
        config_layout.addWidget(self.save_path_input, 1, 1)
        
        self.browse_path_btn = QPushButton("📁")
        self.browse_path_btn.setMaximumWidth(40)
        self.browse_path_btn.clicked.connect(self.browse_save_path)
        self.browse_path_btn.setToolTip("浏览文件夹")
        config_layout.addWidget(self.browse_path_btn, 1, 2)
        
        main_layout.addWidget(config_group)
        
        # 操作区域 - 横向布局，包含按钮和预览
        operation_layout = QHBoxLayout()
        
        # 左侧：按钮和图片列表
        left_panel = QGroupBox("操作")
        left_layout = QVBoxLayout(left_panel)
        left_layout.setSpacing(5)
        
        # 按钮组
        btn_layout = QGridLayout()
        btn_layout.setSpacing(5)
        
        self.screenshot_btn = QPushButton("📷 截图")
        self.screenshot_btn.clicked.connect(self.start_screenshot)
        self.screenshot_btn.setStyleSheet("QPushButton { background-color: #4CAF50; color: white; font-weight: bold; padding: 6px; }")
        btn_layout.addWidget(self.screenshot_btn, 0, 0)
        
        self.select_images_btn = QPushButton("📁 选择")
        self.select_images_btn.clicked.connect(self.select_images)
        btn_layout.addWidget(self.select_images_btn, 0, 1)
        
        self.clear_images_btn = QPushButton("🗑️ 清空")
        self.clear_images_btn.clicked.connect(self.clear_images)
        btn_layout.addWidget(self.clear_images_btn, 1, 0)
        
        self.extract_btn = QPushButton("⚡ 提取")
        self.extract_btn.clicked.connect(self.start_extraction)
        self.extract_btn.setStyleSheet("QPushButton { background-color: #2196F3; color: white; font-weight: bold; padding: 6px; }")
        btn_layout.addWidget(self.extract_btn, 1, 1)
        
        left_layout.addLayout(btn_layout)
        
        # 提取方法显示
        method_layout = QHBoxLayout()
        method_layout.addWidget(QLabel("方法:"))
        self.method_label = QLabel("未选择")
        self.method_label.setStyleSheet("QLabel { color: #ff6b6b; font-weight: bold; padding: 3px; background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 3px; font-size: 10px; }")
        method_layout.addWidget(self.method_label)
        method_layout.addStretch()
        left_layout.addLayout(method_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        self.progress_bar.setMaximumHeight(15)
        left_layout.addWidget(self.progress_bar)
        
        # 图片列表
        self.image_list = QListWidget()
        self.image_list.setMaximumHeight(80)
        self.image_list.setToolTip("图片列表（右键可删除单个图片，双击预览）")
        self.image_list.setContextMenuPolicy(Qt.CustomContextMenu)
        self.image_list.customContextMenuRequested.connect(self.show_image_list_context_menu)
        self.image_list.itemDoubleClicked.connect(self.on_image_list_double_click)
        left_layout.addWidget(self.image_list)
        
        operation_layout.addWidget(left_panel, 1)
        
        # 右侧：预览
        preview_panel = QGroupBox("预览")
        preview_layout = QVBoxLayout(preview_panel)
        preview_layout.setContentsMargins(5, 5, 5, 5)
        
        self.preview_label = QLabel("暂无预览")
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.setMinimumSize(200, 150)
        self.preview_label.setMaximumSize(250, 180)
        self.preview_label.setStyleSheet("border: 2px dashed #ccc; background-color: #f9f9f9;")
        self.preview_label.setScaledContents(False)
        preview_layout.addWidget(self.preview_label)
        
        operation_layout.addWidget(preview_panel, 1)
        
        main_layout.addLayout(operation_layout)
        
        # 结果显示区域
        result_group = QGroupBox("结果")
        result_layout = QVBoxLayout(result_group)
        result_layout.setSpacing(5)
        result_layout.setContentsMargins(5, 5, 5, 5)
        
        # 结果表格
        self.result_table = QTableWidget()
        self.result_table.setColumnCount(6)
        self.result_table.setHorizontalHeaderLabels(["图片名称", "提取方法", "解析方式", "区服", "角色名", "UID"])
        self.result_table.setMaximumHeight(150)
        
        # 设置表格属性
        header = self.result_table.horizontalHeader()
        header.setSectionResizeMode(QHeaderView.Stretch)
        self.result_table.verticalHeader().setDefaultSectionSize(25)  # 减小行高
        
        # 启用表格编辑
        self.result_table.setEditTriggers(QTableWidget.DoubleClicked | QTableWidget.EditKeyPressed)
        
        # 连接单元格编辑信号
        self.result_table.cellChanged.connect(self.on_cell_changed)
        
        # 启用右键菜单
        self.result_table.setContextMenuPolicy(Qt.CustomContextMenu)
        self.result_table.customContextMenuRequested.connect(self.show_context_menu)
        
        result_layout.addWidget(self.result_table)
        
        # 导出按钮和区服提示
        bottom_layout = QHBoxLayout()
        
        # 区服提示（紧凑版）
        tip_text = "区服：铭心之界/拉普达/轻云之月/应许之地/旷野之原/遗忘之境/参天之林/生命之树/魔女小樱/绽放之花/暗涌之泉/缥缈之峰/无尽之海（可输入简称）"
        self.server_tip_label = QLabel(tip_text)
        self.server_tip_label.setStyleSheet("QLabel { background-color: #f8f9fa; padding: 5px; border: 1px solid #dee2e6; border-radius: 3px; font-size: 9px; color: #666; }")
        self.server_tip_label.setWordWrap(True)
        self.server_tip_label.setMaximumHeight(40)
        bottom_layout.addWidget(self.server_tip_label, 2)
        
        self.export_btn = QPushButton("📤 导出")
        self.export_btn.clicked.connect(self.export_to_excel)
        self.export_btn.setEnabled(False)
        self.export_btn.setStyleSheet("QPushButton { background-color: #FF9800; color: white; font-weight: bold; padding: 6px; }")
        self.export_btn.setMaximumWidth(80)
        bottom_layout.addWidget(self.export_btn)
        
        self.import_btn = QPushButton("📥 导入")
        self.import_btn.clicked.connect(self.import_from_excel)
        self.import_btn.setStyleSheet("QPushButton { background-color: #9C27B0; color: white; font-weight: bold; padding: 6px; }")
        self.import_btn.setMaximumWidth(80)
        bottom_layout.addWidget(self.import_btn)
        
        result_layout.addLayout(bottom_layout)
        main_layout.addWidget(result_group)
        
        # 随机抽取区域
        lottery_group = QGroupBox("随机抽取")
        lottery_layout = QHBoxLayout(lottery_group)
        lottery_layout.setSpacing(5)
        lottery_layout.setContentsMargins(5, 5, 5, 5)
        
        lottery_layout.addWidget(QLabel("抽取人数:"))
        
        self.lottery_count_spin = QSpinBox()
        self.lottery_count_spin.setMinimum(1)
        self.lottery_count_spin.setMaximum(100)
        self.lottery_count_spin.setValue(1)
        self.lottery_count_spin.setMaximumWidth(80)
        lottery_layout.addWidget(self.lottery_count_spin)
        
        self.lottery_btn = QPushButton("🎲 开始抽取")
        self.lottery_btn.clicked.connect(self.start_lottery)
        self.lottery_btn.setStyleSheet("QPushButton { background-color: #E91E63; color: white; font-weight: bold; padding: 6px; }")
        self.lottery_btn.setMaximumWidth(100)
        lottery_layout.addWidget(self.lottery_btn)
        
        self.export_lottery_btn = QPushButton("📤 导出")
        self.export_lottery_btn.clicked.connect(self.export_lottery_results)
        self.export_lottery_btn.setEnabled(False)
        self.export_lottery_btn.setStyleSheet("QPushButton { background-color: #4CAF50; color: white; font-weight: bold; padding: 6px; }")
        self.export_lottery_btn.setMaximumWidth(80)
        self.export_lottery_btn.setToolTip("导出抽取结果到Excel")
        lottery_layout.addWidget(self.export_lottery_btn)
        
        lottery_layout.addWidget(QLabel("结果:"))
        
        self.lottery_result_label = QLabel("暂无抽取结果")
        self.lottery_result_label.setStyleSheet("QLabel { background-color: #fff3cd; padding: 5px; border: 1px solid #ffc107; border-radius: 3px; font-weight: bold; color: #856404; }")
        self.lottery_result_label.setWordWrap(True)
        lottery_layout.addWidget(self.lottery_result_label, 1)
        
        main_layout.addWidget(lottery_group)
        
        # 创建菜单栏
        self.create_menu_bar()
        
        # 状态栏
        self.statusBar().showMessage("就绪")
        
        # 加载保存的配置
        self.load_saved_config()
    
    def create_menu_bar(self):
        """创建菜单栏"""
        menubar = self.menuBar()
        
        # 帮助菜单
        help_menu = menubar.addMenu('帮助')
        
        # 联系作者
        contact_action = QAction('联系作者', self)
        contact_action.triggered.connect(self.show_contact_info)
        help_menu.addAction(contact_action)
        
        # 关于
        about_action = QAction('关于', self)
        about_action.triggered.connect(self.show_about)
        help_menu.addAction(about_action)
    
    def show_contact_info(self):
        """显示联系信息"""
        contact_info = """
联系作者信息：

📧 邮箱：AJower@163.com
💬 QQ：3175268324

如有问题或建议，欢迎联系！
        """
        QMessageBox.information(self, "联系作者", contact_info)
    
    def show_about(self):
        """显示关于信息"""
        about_info = """
角色信息提取工具V1.0.0

功能特点：
• 支持截图和文件选择两种方式
• 智能识别区服名称（支持简称）
• AI自动提取角色信息
• 支持手动编辑结果
• 导出Excel表格

技术栈：
• Python + PyQt5
• OpenCV图像处理
• 硅基流动AI模型

作者：AJower
邮箱：AJower@163.com
QQ：3175268324
        """
        QMessageBox.about(self, "关于", about_info)
    
    def load_saved_config(self):
        """加载保存的配置"""
        # 加载API密钥
        if self.api_key:
            self.api_key_input.setText(self.api_key)
        
        # 加载保存路径
        self.save_path_input.setText(self.save_path)
        
        # 更新提取方法显示
        self.update_extraction_method(self.extraction_method)
    
    def closeEvent(self, event):
        """窗口关闭事件"""
        # 保存窗口几何信息
        geometry = {
            "x": self.x(),
            "y": self.y(),
            "width": self.width(),
            "height": self.height()
        }
        self.config_manager.set_window_geometry(geometry)
        
        # 保存当前设置
        self.config_manager.set_api_key(self.api_key_input.text())
        self.config_manager.set_save_path(self.save_path_input.text())
        self.config_manager.set_extraction_method(self.extraction_method)
        
        event.accept()
    
    def browse_save_path(self):
        """浏览保存路径"""
        folder = QFileDialog.getExistingDirectory(self, "选择截图保存路径", self.save_path)
        if folder:
            self.save_path = folder
            self.save_path_input.setText(folder)
            self.statusBar().showMessage(f"保存路径已设置为: {folder}")
    
    def update_extraction_method(self, method):
        """更新提取方法显示"""
        self.extraction_method = method
        if method == "screenshot":
            self.method_label.setText("📷 截图")
            self.method_label.setStyleSheet("QLabel { color: #28a745; font-weight: bold; padding: 3px; background-color: #d4edda; border: 1px solid #c3e6cb; border-radius: 3px; font-size: 10px; }")
        elif method == "file":
            self.method_label.setText("📁 文件")
            self.method_label.setStyleSheet("QLabel { color: #007bff; font-weight: bold; padding: 3px; background-color: #d1ecf1; border: 1px solid #bee5eb; border-radius: 3px; font-size: 10px; }")
        else:
            self.method_label.setText("未选择")
            self.method_label.setStyleSheet("QLabel { color: #ff6b6b; font-weight: bold; padding: 3px; background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 3px; font-size: 10px; }")
    
    def on_cell_changed(self, row, column):
        """表格单元格编辑事件"""
        if row < len(self.results):
            item = self.result_table.item(row, column)
            if item:
                new_value = item.text()
                # 调整列索引，因为添加了提取方法和解析方式列
                if column == 0:  # 图片名称列
                    self.results[row]['image_path'] = new_value
                elif column == 3:  # 区服列
                    # 智能匹配区服名称
                    matched_server = self.match_server_name(new_value)
                    if matched_server != new_value:
                        # 如果匹配到完整名称，更新显示
                        self.results[row]['server'] = matched_server
                        self.result_table.setItem(row, column, QTableWidgetItem(matched_server))
                        self.statusBar().showMessage(f"已更新第{row+1}行第{column+1}列: {new_value} → {matched_server}")
                    else:
                        self.results[row]['server'] = new_value
                        self.statusBar().showMessage(f"已更新第{row+1}行第{column+1}列: {new_value}")
                elif column == 4:  # 角色名列
                    self.results[row]['character_name'] = new_value
                    self.statusBar().showMessage(f"已更新第{row+1}行第{column+1}列: {new_value}")
                elif column == 5:  # UID列
                    self.results[row]['uid'] = new_value
                    self.statusBar().showMessage(f"已更新第{row+1}行第{column+1}列: {new_value}")
                # 提取方法列（第1列）和解析方式列（第2列）不允许编辑
    
    def match_server_name(self, input_name):
        """智能匹配区服名称"""
        if not input_name or not input_name.strip():
            return input_name
        
        input_name = input_name.strip()
        
        # 如果输入的是完整名称，直接返回
        if input_name in self.server_names:
            return input_name
        
        # 创建区服名称的简称映射
        server_shortcuts = {
            # 基于前几个字符的匹配
            "铭心": "铭心之界",
            "拉普": "拉普达", 
            "轻云": "轻云之月",
            "应许": "应许之地",
            "旷野": "旷野之原",
            "遗忘": "遗忘之境",
            "参天": "参天之林",
            "生命": "生命之树",
            "魔女": "魔女小樱",
            "绽放": "绽放之花",
            "暗涌": "暗涌之泉",
            "缥缈": "缥缈之峰",
            "无尽": "无尽之海",
            
            # 基于关键词的匹配
            "界": "铭心之界",
            "达": "拉普达",
            "月": "轻云之月", 
            "地": "应许之地",
            "原": "旷野之原",
            "境": "遗忘之境",
            "林": "参天之林",
            "树": "生命之树",
            "樱": "魔女小樱",
            "花": "绽放之花",
            "泉": "暗涌之泉",
            "峰": "缥缈之峰",
            "海": "无尽之海"
        }
        
        # 首先尝试精确匹配
        if input_name in server_shortcuts:
            return server_shortcuts[input_name]
        
        # 尝试模糊匹配
        for shortcut, full_name in server_shortcuts.items():
            if shortcut in input_name or input_name in shortcut:
                return full_name
        
        # 尝试包含匹配
        for server_name in self.server_names:
            if input_name in server_name or server_name in input_name:
                return server_name
        
        # 如果没有匹配到，返回原输入
        return input_name
    
    def show_context_menu(self, position):
        """显示右键菜单"""
        if self.result_table.itemAt(position) is None:
            return
        
        # 获取当前行
        current_row = self.result_table.rowAt(position.y())
        if current_row < 0 or current_row >= len(self.results):
            return
        
        # 创建右键菜单
        context_menu = QMenu(self)
        
        # 删除行动作
        delete_action = QAction("🗑️ 删除此行", self)
        delete_action.triggered.connect(lambda: self.delete_row(current_row))
        context_menu.addAction(delete_action)
        
        # 显示菜单
        context_menu.exec_(self.result_table.mapToGlobal(position))
    
    def delete_row(self, row):
        """删除指定行"""
        if 0 <= row < len(self.results):
            # 确认删除
            reply = QMessageBox.question(self, "确认删除", 
                                       f"确定要删除第{row+1}行的数据吗？",
                                       QMessageBox.Yes | QMessageBox.No)
            
            if reply == QMessageBox.Yes:
                # 删除结果数据
                del self.results[row]
                
                # 重新显示表格
                self.display_results()
                
                # 更新状态
                self.statusBar().showMessage(f"已删除第{row+1}行数据")
                
                # 如果没有数据了，禁用导出按钮
                if not self.results:
                    self.export_btn.setEnabled(False)
    
    def start_screenshot(self):
        """开始截图"""
        try:
            print("开始截图流程...")
            # 隐藏主窗口
            self.hide()
            print("主窗口已隐藏")
            
            # 等待一小段时间确保窗口完全隐藏
            QTimer.singleShot(200, self.show_screenshot_widget)
            
        except Exception as e:
            print(f"启动截图失败: {e}")
            QMessageBox.critical(self, "错误", f"启动截图失败：{str(e)}")
            self.show()
    
    def show_screenshot_widget(self):
        """显示截图小部件"""
        try:
            print("创建截图小部件...")
            # 更新保存路径
            self.save_path = self.save_path_input.text() or os.getcwd()
            self.screenshot_widget = ScreenshotWidget(self.save_path)
            self.screenshot_widget.screenshot_taken.connect(self.on_screenshot_taken)
            self.screenshot_widget.start_screenshot()
        except Exception as e:
            print(f"截图失败: {e}")
            import traceback
            traceback.print_exc()
            QMessageBox.critical(self, "错误", f"截图失败：{str(e)}")
            self.show()
    
    def on_screenshot_taken(self, image_path):
        """截图完成回调"""
        try:
            # 显示主窗口
            self.show()
            self.raise_()
            self.activateWindow()
            
            # 添加到图片列表
            self.image_paths.append(image_path)
            self.update_image_list()
            
            # 更新预览
            self.update_preview(image_path)
            
            # 更新提取方法
            self.update_extraction_method("screenshot")
            
            # 更新状态
            self.statusBar().showMessage(f"截图完成！已添加 {os.path.basename(image_path)}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"处理截图失败：{str(e)}")
    
    def update_image_list(self):
        """更新图片列表显示"""
        self.image_list.clear()
        if self.image_paths:
            for path in self.image_paths:
                image_name = os.path.basename(path)
                item = QListWidgetItem(f"📷 {image_name}")
                item.setData(Qt.UserRole, path)  # 存储完整路径
                item.setToolTip(path)  # 鼠标悬停显示完整路径
                self.image_list.addItem(item)
    
    def update_preview(self, image_path):
        """更新截图预览"""
        try:
            pixmap = QPixmap(image_path)
            if not pixmap.isNull():
                # 缩放图片以适应预览区域（适配紧凑布局）
                scaled_pixmap = pixmap.scaled(240, 170, Qt.KeepAspectRatio, Qt.SmoothTransformation)
                self.preview_label.setPixmap(scaled_pixmap)
                self.preview_label.setText("")
            else:
                self.preview_label.setText("预览加载失败")
        except Exception as e:
            self.preview_label.setText(f"预览错误：{str(e)}")
    
    def show_image_list_context_menu(self, position):
        """显示图片列表右键菜单"""
        item = self.image_list.itemAt(position)
        if item is None:
            return
        
        # 创建右键菜单
        context_menu = QMenu(self)
        
        # 预览动作
        preview_action = QAction("👁️ 预览", self)
        preview_action.triggered.connect(lambda: self.on_image_list_double_click(item))
        context_menu.addAction(preview_action)
        
        # 删除动作
        delete_action = QAction("🗑️ 删除", self)
        delete_action.triggered.connect(lambda: self.delete_image_from_list(item))
        context_menu.addAction(delete_action)
        
        # 显示菜单
        context_menu.exec_(self.image_list.mapToGlobal(position))
    
    def on_image_list_double_click(self, item):
        """双击图片列表项时预览"""
        if item:
            image_path = item.data(Qt.UserRole)
            self.update_preview(image_path)
    
    def delete_image_from_list(self, item):
        """从列表中删除指定图片"""
        if not item:
            return
        
        image_path = item.data(Qt.UserRole)
        image_name = os.path.basename(image_path)
        
        # 确认删除
        reply = QMessageBox.question(self, "确认删除", 
                                   f"确定要从列表中删除 {image_name} 吗？",
                                   QMessageBox.Yes | QMessageBox.No)
        
        if reply == QMessageBox.Yes:
            try:
                # 从图片路径列表中删除
                if image_path in self.image_paths:
                    self.image_paths.remove(image_path)
                
                # 更新列表显示
                self.update_image_list()
                
                # 清除预览（如果删除的是当前预览的图片）
                current_pixmap = self.preview_label.pixmap()
                if current_pixmap and not current_pixmap.isNull():
                    # 如果列表为空，清除预览
                    if not self.image_paths:
                        self.preview_label.setText("暂无预览")
                        self.preview_label.setPixmap(QPixmap())
                    # 否则显示第一张图片
                    else:
                        self.update_preview(self.image_paths[0])
                
                # 更新状态
                self.statusBar().showMessage(f"已删除 {image_name}")
                
                # 如果列表为空，重置提取方法
                if not self.image_paths:
                    self.update_extraction_method("未选择")
                    
            except Exception as e:
                QMessageBox.critical(self, "错误", f"删除失败：{str(e)}")
    
    def select_images(self):
        """选择图片文件"""
        file_dialog = QFileDialog()
        file_dialog.setFileMode(QFileDialog.ExistingFiles)
        file_dialog.setNameFilter("图片文件 (*.png *.jpg *.jpeg *.bmp *.gif)")
        
        if file_dialog.exec_():
            self.image_paths = file_dialog.selectedFiles()
            self.update_image_list()
            
            # 更新提取方法
            self.update_extraction_method("file")
            
            # 更新预览（显示第一张图片）
            if self.image_paths:
                self.update_preview(self.image_paths[0])
            
            self.statusBar().showMessage(f"已选择 {len(self.image_paths)} 张图片")
    
    def clear_images(self):
        """清空图片列表"""
        # 删除已保存的截图文件
        for image_path in self.image_paths:
            try:
                if os.path.exists(image_path) and image_path.startswith(os.getcwd()):
                    os.remove(image_path)
            except Exception as e:
                print(f"删除文件失败 {image_path}: {e}")
        
        self.image_paths = []
        self.image_list.clear()
        self.result_table.setRowCount(0)
        self.export_btn.setEnabled(False)
        self.preview_label.setText("暂无截图预览")
        self.preview_label.setPixmap(QPixmap())
        
        # 重置提取方法
        self.update_extraction_method("未选择")
        
        self.statusBar().showMessage("已清空图片列表")
    
    def start_extraction(self):
        """开始提取信息"""
        if not self.api_key_input.text().strip():
            QMessageBox.warning(self, "警告", "请输入API密钥！")
            return
        
        if not self.image_paths:
            QMessageBox.warning(self, "警告", "请先选择图片文件！")
            return
        
        self.api_key = self.api_key_input.text().strip()
        
        # 禁用按钮
        self.extract_btn.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # 创建并启动提取线程
        self.extractor_thread = AIExtractor(self.api_key, self.image_paths, self.server_names, self.extraction_method)
        self.extractor_thread.progress_updated.connect(self.update_progress)
        self.extractor_thread.result_ready.connect(self.extraction_completed)
        self.extractor_thread.error_occurred.connect(self.extraction_error)
        self.extractor_thread.start()
        
        self.statusBar().showMessage("正在提取信息...")
    
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
    
    def extraction_completed(self, data):
        """提取完成"""
        self.results = data['results']
        self.display_results()
        
        # 恢复界面状态
        self.extract_btn.setEnabled(True)
        self.progress_bar.setVisible(False)
        self.export_btn.setEnabled(True)
        
        self.statusBar().showMessage(f"提取完成！共处理 {len(self.results)} 张图片")
    
    def extraction_error(self, error_msg):
        """提取出错"""
        QMessageBox.critical(self, "错误", f"提取过程中出现错误：{error_msg}")
        
        # 恢复界面状态
        self.extract_btn.setEnabled(True)
        self.progress_bar.setVisible(False)
        self.statusBar().showMessage("提取失败")
    
    def display_results(self):
        """显示提取结果"""
        self.result_table.setRowCount(len(self.results))
        
        for i, result in enumerate(self.results):
            self.result_table.setItem(i, 0, QTableWidgetItem(result['image_path']))
            
            # 显示提取方法
            method_text = "📷 截图" if result.get('method') == 'screenshot' else "📁 文件"
            self.result_table.setItem(i, 1, QTableWidgetItem(method_text))
            
            # 显示解析方式
            parse_type = result.get('parse_type', '未知')
            self.result_table.setItem(i, 2, QTableWidgetItem(parse_type))
            
            self.result_table.setItem(i, 3, QTableWidgetItem(result['server']))
            self.result_table.setItem(i, 4, QTableWidgetItem(result['character_name']))
            self.result_table.setItem(i, 5, QTableWidgetItem(result['uid']))
    
    def import_from_excel(self):
        """从Excel文件导入数据"""
        file_dialog = QFileDialog()
        file_dialog.setFileMode(QFileDialog.ExistingFile)
        file_dialog.setNameFilter("Excel文件 (*.xlsx *.xls)")
        
        if file_dialog.exec_():
            file_path = file_dialog.selectedFiles()[0]
            
            try:
                # 读取Excel文件
                df = pd.read_excel(file_path, engine='openpyxl')
                
                # 检查必需的列
                required_columns = ['区服', '角色名', 'UID']
                missing_columns = [col for col in required_columns if col not in df.columns]
                
                if missing_columns:
                    QMessageBox.warning(self, "导入失败", 
                                      f"Excel文件缺少必需的列: {', '.join(missing_columns)}\n"
                                      f"需要包含: {', '.join(required_columns)}")
                    return
                
                # 转换为results格式
                self.results = []
                for index, row in df.iterrows():
                    result = {
                        'image_path': f"导入数据_{index + 1}",
                        'method': 'import',
                        'parse_type': '导入',
                        'server': str(row.get('区服', '未知')),
                        'character_name': str(row.get('角色名', '未知')),
                        'uid': str(row.get('UID', '未知')),
                        'error': None
                    }
                    self.results.append(result)
                
                # 显示结果
                self.display_results()
                
                # 启用导出按钮
                self.export_btn.setEnabled(True)
                
                # 更新状态
                QMessageBox.information(self, "导入成功", 
                                      f"成功导入 {len(self.results)} 条数据！\n"
                                      f"现在可以进行随机抽取了。")
                self.statusBar().showMessage(f"已导入 {len(self.results)} 条数据")
                
                # 重置抽取结果
                self.lottery_result_label.setText("暂无抽取结果")
                self.lottery_selected_results = []
                self.export_lottery_btn.setEnabled(False)
                
            except Exception as e:
                import traceback
                traceback.print_exc()
                QMessageBox.critical(self, "导入失败", f"导入Excel文件失败：{str(e)}")
    
    def start_lottery(self):
        """开始随机抽取"""
        if not self.results:
            QMessageBox.warning(self, "提示", "请先导入数据或提取信息后再进行抽取！")
            return
        
        lottery_count = self.lottery_count_spin.value()
        
        # 检查抽取人数是否超过总人数
        if lottery_count > len(self.results):
            QMessageBox.warning(self, "提示", 
                              f"抽取人数({lottery_count})不能超过总人数({len(self.results)})！")
            return
        
        try:
            import random
            
            # 随机抽取
            selected_results = random.sample(self.results, lottery_count)
            
            # 保存抽取结果
            self.lottery_selected_results = selected_results
            
            # 格式化结果显示
            result_text = []
            for i, result in enumerate(selected_results, 1):
                char_name = result.get('character_name', '未知')
                server = result.get('server', '未知')
                uid = result.get('uid', '未知')
                result_text.append(f"{i}. {char_name} ({server} - UID:{uid})")
            
            # 显示结果
            result_display = " | ".join(result_text)
            self.lottery_result_label.setText(result_display)
            
            # 在表格中高亮显示抽中的行
            self.highlight_lottery_results(selected_results)
            
            # 启用导出按钮
            self.export_lottery_btn.setEnabled(True)
            
            # 弹窗显示完整结果
            detail_text = "\n".join(result_text)
            QMessageBox.information(self, "🎉 抽取结果", 
                                  f"恭喜以下 {lottery_count} 位幸运儿：\n\n{detail_text}")
            
            self.statusBar().showMessage(f"已随机抽取 {lottery_count} 人，可以导出结果")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"抽取失败：{str(e)}")
    
    def highlight_lottery_results(self, selected_results):
        """在表格中高亮显示抽取结果"""
        try:
            # 先清除所有高亮
            for row in range(self.result_table.rowCount()):
                for col in range(self.result_table.columnCount()):
                    item = self.result_table.item(row, col)
                    if item:
                        item.setBackground(QColor(255, 255, 255))  # 白色背景
            
            # 高亮抽中的行
            for selected in selected_results:
                for row in range(self.result_table.rowCount()):
                    table_char_name = self.result_table.item(row, 4)  # 角色名列
                    table_uid = self.result_table.item(row, 5)  # UID列
                    
                    if (table_char_name and table_uid and 
                        table_char_name.text() == selected.get('character_name') and
                        table_uid.text() == selected.get('uid')):
                        # 设置为黄色高亮
                        for col in range(self.result_table.columnCount()):
                            item = self.result_table.item(row, col)
                            if item:
                                item.setBackground(QColor(255, 255, 0, 100))  # 半透明黄色
                        break
        except Exception as e:
            print(f"高亮显示失败: {e}")
    
    def export_to_excel(self):
        """导出到Excel文件"""
        if not self.results:
            QMessageBox.warning(self, "警告", "没有可导出的数据！")
            return
        
        file_dialog = QFileDialog()
        file_dialog.setAcceptMode(QFileDialog.AcceptSave)
        file_dialog.setNameFilter("Excel文件 (*.xlsx)")
        file_dialog.setDefaultSuffix("xlsx")
        
        if file_dialog.exec_():
            file_path = file_dialog.selectedFiles()[0]
            
            try:
                # 创建DataFrame，只保留需要的列
                df = pd.DataFrame(self.results)
                
                # 只保留区服、角色名、UID三列
                df = df[['server', 'character_name', 'uid']]
                
                # 重命名列以便更好地显示
                column_mapping = {
                    'server': '区服',
                    'character_name': '角色名',
                    'uid': 'UID'
                }
                df = df.rename(columns=column_mapping)
                
                # 导出到Excel
                df.to_excel(file_path, index=False, engine='openpyxl')
                
                # 添加图片到Excel
                from openpyxl import load_workbook
                from openpyxl.drawing.image import Image as ExcelImage
                
                wb = load_workbook(file_path)
                ws = wb.active
                
                # 在第一列之前插入一列用于显示图片
                ws.insert_cols(1)
                ws.cell(row=1, column=1, value="截图预览")
                
                # 设置列宽和行高（统一格式）
                ws.column_dimensions['A'].width = 15  # 图片列宽度（缩小）
                
                # 固定的图片尺寸（统一格式）
                fixed_img_width = 80   # 固定宽度（像素）
                fixed_img_height = 60  # 固定高度（像素）
                
                # 遍历每个结果，添加图片
                for i, result in enumerate(self.results):
                    row_idx = i + 2  # Excel行从1开始，第1行是标题
                    
                    # 设置统一的行高
                    ws.row_dimensions[row_idx].height = 50  # 缩小行高
                    
                    # 获取图片路径
                    image_name = result['image_path']
                    
                    # 尝试找到完整路径
                    image_full_path = None
                    
                    # 如果是截图提取，图片在保存路径中
                    if result.get('method') == 'screenshot':
                        possible_path = os.path.join(self.save_path, image_name)
                        if os.path.exists(possible_path):
                            image_full_path = possible_path
                    
                    # 如果是文件提取，从原始路径列表中查找
                    if not image_full_path:
                        for img_path in self.image_paths:
                            if os.path.basename(img_path) == image_name or img_path == image_name:
                                image_full_path = img_path
                                break
                    
                    # 如果直接是完整路径
                    if not image_full_path and os.path.exists(image_name):
                        image_full_path = image_name
                    
                    # 插入图片
                    if image_full_path and os.path.exists(image_full_path):
                        try:
                            # 创建Excel图片对象
                            img = ExcelImage(image_full_path)
                            
                            # 统一设置为固定尺寸（所有图片都是相同大小）
                            img.width = fixed_img_width
                            img.height = fixed_img_height
                            
                            # 设置图片锚点（插入到单元格）
                            cell_address = f'A{row_idx}'
                            img.anchor = cell_address
                            
                            # 添加图片到工作表
                            ws.add_image(img)
                            
                        except Exception as e:
                            print(f"插入图片失败 {image_full_path}: {e}")
                            ws.cell(row=row_idx, column=1, value="图片加载失败")
                    else:
                        ws.cell(row=row_idx, column=1, value="图片未找到")
                
                # 保存工作簿
                wb.save(file_path)
                
                QMessageBox.information(self, "成功", f"数据和图片已导出到：{file_path}")
                self.statusBar().showMessage(f"数据和图片已导出到：{file_path}")
                
            except Exception as e:
                import traceback
                traceback.print_exc()
                QMessageBox.critical(self, "错误", f"导出失败：{str(e)}")
    
    def export_lottery_results(self):
        """导出抽取结果到Excel文件"""
        if not self.lottery_selected_results:
            QMessageBox.warning(self, "警告", "没有抽取结果可以导出！\n请先进行随机抽取。")
            return
        
        file_dialog = QFileDialog()
        file_dialog.setAcceptMode(QFileDialog.AcceptSave)
        file_dialog.setNameFilter("Excel文件 (*.xlsx)")
        file_dialog.setDefaultSuffix("xlsx")
        file_dialog.selectFile("抽取结果.xlsx")  # 默认文件名
        
        if file_dialog.exec_():
            file_path = file_dialog.selectedFiles()[0]
            
            try:
                # 创建DataFrame，只保留需要的列
                df = pd.DataFrame(self.lottery_selected_results)
                
                # 只保留区服、角色名、UID三列
                df = df[['server', 'character_name', 'uid']]
                
                # 重命名列以便更好地显示
                column_mapping = {
                    'server': '区服',
                    'character_name': '角色名',
                    'uid': 'UID'
                }
                df = df.rename(columns=column_mapping)
                
                # 添加序号列
                df.insert(0, '序号', range(1, len(df) + 1))
                
                # 导出到Excel
                df.to_excel(file_path, index=False, engine='openpyxl')
                
                # 添加图片到Excel
                from openpyxl import load_workbook
                from openpyxl.drawing.image import Image as ExcelImage
                
                wb = load_workbook(file_path)
                ws = wb.active
                
                # 在第二列之前插入一列用于显示图片（序号列后面）
                ws.insert_cols(2)
                ws.cell(row=1, column=2, value="截图预览")
                
                # 设置列宽和行高（统一格式）
                ws.column_dimensions['B'].width = 15  # 图片列宽度
                
                # 固定的图片尺寸（统一格式）
                fixed_img_width = 80   # 固定宽度（像素）
                fixed_img_height = 60  # 固定高度（像素）
                
                # 遍历每个抽取结果，添加图片
                for i, result in enumerate(self.lottery_selected_results):
                    row_idx = i + 2  # Excel行从1开始，第1行是标题
                    
                    # 设置统一的行高
                    ws.row_dimensions[row_idx].height = 50
                    
                    # 获取图片路径
                    image_name = result.get('image_path', '')
                    
                    # 尝试找到完整路径（只有非导入数据才有图片）
                    image_full_path = None
                    
                    if not image_name.startswith('导入数据_'):
                        # 如果是截图提取，图片在保存路径中
                        if result.get('method') == 'screenshot':
                            possible_path = os.path.join(self.save_path, image_name)
                            if os.path.exists(possible_path):
                                image_full_path = possible_path
                        
                        # 如果是文件提取，从原始路径列表中查找
                        if not image_full_path:
                            for img_path in self.image_paths:
                                if os.path.basename(img_path) == image_name or img_path == image_name:
                                    image_full_path = img_path
                                    break
                        
                        # 如果直接是完整路径
                        if not image_full_path and os.path.exists(image_name):
                            image_full_path = image_name
                    
                    # 插入图片
                    if image_full_path and os.path.exists(image_full_path):
                        try:
                            # 创建Excel图片对象
                            img = ExcelImage(image_full_path)
                            
                            # 统一设置为固定尺寸
                            img.width = fixed_img_width
                            img.height = fixed_img_height
                            
                            # 设置图片锚点
                            cell_address = f'B{row_idx}'
                            img.anchor = cell_address
                            
                            # 添加图片到工作表
                            ws.add_image(img)
                            
                        except Exception as e:
                            print(f"插入图片失败 {image_full_path}: {e}")
                            ws.cell(row=row_idx, column=2, value="图片加载失败")
                    else:
                        ws.cell(row=row_idx, column=2, value="无图片")
                
                # 保存工作簿
                wb.save(file_path)
                
                QMessageBox.information(self, "导出成功", 
                                      f"抽取结果已导出到：{file_path}\n"
                                      f"共 {len(self.lottery_selected_results)} 位幸运儿！")
                self.statusBar().showMessage(f"抽取结果已导出到：{file_path}")
                
            except Exception as e:
                import traceback
                traceback.print_exc()
                QMessageBox.critical(self, "错误", f"导出失败：{str(e)}")


def main():
    """主函数"""
    print("=" * 60)
    print("角色信息提取工具V1.0.0")
    print("作者：AJower")
    print("邮箱：AJower@163.com")
    print("QQ：3175268324")
    print("=" * 60)
    
    # 检查并安装依赖
    try:
        missing_deps = check_dependencies()
        if missing_deps:
            print(f"发现缺失的依赖包: {missing_deps}")
            print("正在自动安装依赖...")
            if not install_dependencies():
                print("依赖安装失败，程序可能无法正常运行")
                input("按回车键退出...")
                return
    except Exception as e:
        print(f"依赖检查失败: {e}")
        input("按回车键退出...")
        return
    
    app = QApplication(sys.argv)
    app.setApplicationName("角色信息提取工具V1.0.0")
    app.setApplicationVersion("1.0.0")
    
    # 设置应用样式
    app.setStyle('Fusion')
    
    # 创建登录对话框
    login_dialog = LoginDialog()
    
    # 显示登录对话框
    if login_dialog.exec_() == QDialog.Accepted:
        # 登录成功，启动主程序
        window = ScreenshotExtractor()
        window.show()
        sys.exit(app.exec_())
    else:
        # 登录失败或取消，退出
        sys.exit(0)


if __name__ == '__main__':
    main()
